import time
from types import FunctionType

import gym
import ril_env

import keyboard
from ril_env.envs.grid_world_v1 import GridWorldEnv
import argparse

action, pause = 0, False


def mcc_preprocess():
    global action
    action = [.0, .0]


def mcc_human_play(key):
    global action, pause
    if key.event_type == 'down' and key.name == 'left':
        action[0] = -1
    elif key.event_type == 'down' and key.name == 'right':
        action[0] = 1


def grid_world_human_play(key):
    global action, pause
    if key.event_type == 'down' and key.name == 'up':
        action = GridWorldEnv.MOVE_UP
    elif key.event_type == 'down' and key.name == 'down':
        action = GridWorldEnv.MOVE_DOWN
    elif key.event_type == 'down' and key.name == 'left':
        action = GridWorldEnv.MOVE_LEFT
    elif key.event_type == 'down' and key.name == 'right':
        action = GridWorldEnv.MOVE_RIGHT


def cart_pole_human_play(key):
    global action, pause
    if key.event_type == 'down' and key.name == 'left':
        action = 0
    elif key.event_type == 'down' and key.name == 'right':
        action = 1


def flappy_bird_human_play(key):
    global action, pause
    if key.event_type == 'down' and key.name == 'space':
        action = 1
    else:
        action = 0


game_factory = {
    'GridWorld-v0': grid_world_human_play,
    'GridWorld-v1': grid_world_human_play,
    'CartPole-v0': cart_pole_human_play,
    'CartPole-v1': cart_pole_human_play,
    'MountainCarContinuous-v0': {
        'preprocess': mcc_preprocess,
        'key_handle': mcc_human_play
    },
    'FlappyBird-v1': flappy_bird_human_play,
}


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--task', type=str, default='GridWorld-v1')
    parser.add_argument('--idle', type=float, default=0.05)
    args = parser.parse_args()
    env = gym.make(args.task)
    env.reset()

    game_fn = game_factory[args.task]
    if isinstance(game_fn, FunctionType):
        keyboard.hook(game_fn)
    else:
        game_fn['preprocess']()
        keyboard.hook(game_fn['key_handle'])

    score = 0
    while True:
        if not pause:
            observation, reward, done, info = env.step(action)
            if done:
                print('Final score: %d' % (score + reward))
                env.close()
                break
            score += reward
        env.render()
        time.sleep(args.idle)
